﻿using System;
using System.Diagnostics;
using System.Linq;
using System.Management;

namespace Toolbar
{
    public class UsbDriveNotificator : IDisposable
    {
        private readonly ManagementEventWatcher m_Watcher;
        public event Action<string> DriveConnected;
        public event Action DriveDisconnected;

        public UsbDriveNotificator()
        {
            const string query = "SELECT * FROM __InstanceOperationEvent WITHIN 1 WHERE TargetInstance ISA \"Win32_DiskDrive\"";
            m_Watcher = new ManagementEventWatcher(query);
            m_Watcher.EventArrived += MediaConnectWatcherEventArrived;
            m_Watcher.Start();
        }

        private void MediaConnectWatcherEventArrived(object sender, EventArrivedEventArgs e)
        {
            try
            {
                var mbo = e.NewEvent;
                Log("arrived: " + mbo.ClassPath.ClassName);
                var obj = (ManagementBaseObject)mbo["TargetInstance"];
                var name = obj["Name"].ToString();
                //var caption = obj["Caption"].ToString();

                switch (mbo.ClassPath.ClassName)
                {
                    case "__InstanceCreationEvent":
                        var driveLetter = GetDriveLetterFromDisk(name);
                        var serial = GetVolumeSerial(driveLetter);
                        if (obj["InterfaceType"].ToString() == "USB")
                        {
                            Log("Insert");
                            FireDriveConnected(serial);
                        }
                        break;

                    case "__InstanceDeletionEvent":

                        if (obj["InterfaceType"].ToString() == "USB")
                        {
                            Log("Out");
                            FireDriveDisconnected();
                        }
                        break;
                }
            }
            catch (Exception exception)
            {
                Log("ex: {0}", exception);
            }
        }

        private static void Log(string format, params object[] args)
        {
            Debug.WriteLine(string.Format(format, args), "UsbDriveNotificator");
        }
        private static string GetVolumeSerial(string strDriveLetter)
        {
            if (string.IsNullOrEmpty(strDriveLetter))
            {
                strDriveLetter = "C:";
            }
            strDriveLetter = strDriveLetter.TrimEnd('\\');
            var disk = new ManagementObject("win32_logicaldisk.deviceid=\"" + strDriveLetter + "\"");
            disk.Get();
            var serial = disk["VolumeSerialNumber"];

            if (serial != null)
            {
                return serial.ToString();
            }

            return null;
        }
        private static string GetDriveLetterFromDisk(string name)
        {
            try
            {
                string ans = string.Empty;

                name = name.Replace(@"\", @"\\");

                var oqPart = new ObjectQuery("ASSOCIATORS OF {Win32_DiskDrive.DeviceID=\"" + name + "\"} WHERE AssocClass = Win32_DiskDriveToDiskPartition");
                var mosPart = new ManagementObjectSearcher(oqPart);

                var collection = mosPart.Get();
                var searchers = (from ManagementObject objPart in collection
                                 select new ObjectQuery("ASSOCIATORS OF {Win32_DiskPartition.DeviceID=\"" + objPart["DeviceID"] + "\"} WHERE AssocClass = Win32_LogicalDiskToPartition")
                                     into queryString
                                     select new ManagementObjectSearcher(queryString));

                Func<string, ManagementObject, string> func1 = (current, disk) => current + (disk["Name"] + ",");
                Func<string, ManagementObjectSearcher, string> func = (current1, mosDisk) => mosDisk.Get().Cast<ManagementObject>().Aggregate(current1, func1);
                ans = searchers.Aggregate(ans, func);

                return ans.Trim(',');
            }
            catch (Exception e)
            {
                Log("GetDriveLetterFromDisk: {0}", e);
                return null;
            }
        }

        public void FireDriveDisconnected()
        {
            var handler = DriveDisconnected;
            if (handler != null)
            {
                handler();
            }
        }

        private void FireDriveConnected(string id)
        {
            var handler = DriveConnected;
            if (handler != null)
            {
                handler(id);
            }
        }

        public void Dispose()
        {
            m_Watcher.Stop();
        }

    }
}
